from autogen import AssistantAgent, ConversableAgent
from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent

from utils.configuration import config, llm_config
from prompts.prompt import *
from tools.tools import fetch_context

# cooperation between agents

memoryAgent = AssistantAgent(name='memory expert',
                             system_message=MemoryAgentPrompt,
                             llm_config=llm_config)
memoryAgent.register_for_llm(name="context_fetcher",
                             description="useful to get the related context for the question")(fetch_context)

cpuAgent = AssistantAgent(name='cpu expert',
                            system_message=CPUAgentPrompt,
                            llm_config=llm_config)
cpuAgent.register_for_llm(name="context_fetcher",
                          description="useful to get the related context for the question")(fetch_context)


summaryAgent = AssistantAgent(name='summary expert',
                            system_message=SummaryAgentPrompt,
                            llm_config=llm_config)


userAgent = ConversableAgent(
    name="User",
    llm_config=False,
    max_consecutive_auto_reply=3,
    # is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
)
userAgent.register_for_execution(name="context_fetcher")(fetch_context)


def cooperation(question):
    res = fetch_context(question)
    expert_template = f"here is context: {res}.  here is my question {question}"
    summary_template = f"My question is: {question} \n \n Here is analysis of other experts: \n"
    memory_result = userAgent.initiate_chat(memoryAgent, message=expert_template,max_turns=1)
    cpu_result = userAgent.initiate_chat(cpuAgent, message=expert_template,max_turns=1)
    for agent in [memoryAgent,cpuAgent]:
        summary_template += f"{agent.name} said: {agent.chat_messages[userAgent][-1]['content']}"
    final = userAgent.initiate_chat(summaryAgent, message=summary_template,max_turns=1)
    print(final)


if __name__ == '__main__':
    # chat_result=userAgent.initiate_chat(memoryAgent,message="为什么hologress会产生内存溢出？需要如何优化")
    # print(chat_result)
    cooperation("为什么hologress会产生内存溢出？需要如何优化")


